import os.path
import re
from typing import Union, List, Dict, Callable, Iterable

from scripts.Log import Log
from scripts.device.AdbDevice import AdbDevice
from scripts.device.AdbShell import PmListPkgInfo, AdbShell
from scripts.util.ProgressHelper import ProgressHelper
from scripts.util.Util import Util
from scripts.actions.Action import Action
from scripts.actions.Context import Context


class BackupAppDataAction(Action):
    KEY = Action.register('app-data')
    DISPLAY_NAME = 'App data'

    EXCLUDE_LIST = {
        'android': None,
        'com.android.backupconfirm': None,
        'com.android.bluetooth': None,
        'com.android.bluetoothmidiservice': None,
        'com.android.BluetoothTestMode': None,
        'com.android.cellbroadcastreceiver': None,
        'com.android.certinstaller': None,
        'com.android.cts.ctsshim': None,
        'com.android.cts.priv.ctsshim': None,
        'com.android.egg': None,
        'com.android.emergency': None,
        'com.android.externalstorage': None,
        'com.android.htmlviewer': None,
        'com.android.mmitest': None,
        'com.android.mtp': None,
        'com.android.musicfx': None,
        'com.android.nfc': None,
        'com.android.NfcTestMode': None,
        'com.android.phone': None,
        'com.android.printspooler': None,
        'com.android.protips': None,
        'com.android.providers.downloads': None,
        # 'com.android.providers.media': None,
        # 'com.android.providers.media.module': None,
        'com.android.server.telecom': None,
        'com.android.shell': None,
        'com.android.soundpicker': None,
        'com.android.statementservice': None,
        'com.android.storagemanager': None,
        'com.android.systemui': None,
        'com.android.systemui.theme.dark': None,
        'com.android.traceur': None,
        'com.android.vpndialogs': None,
        'com.google.android.apps.restore': None,
        'com.google.android.calculator': None,
        'com.google.android.packageinstaller': None,
        'com.google.android.permissioncontroller': None,
        'com.google.android.printservice.recommendation': None,
        'com.google.android.setupwizard': None,
        'com.hidden.engmenu': None,
        'com.qti.ltebc': None,
        'com.qti.qualcomm.datastatusnotification': None,
        'com.qti.qualcomm.deviceinfo': None,
        'com.qti.xdivert': None,
        'com.qualcomm.atfwd': None,
        'com.qualcomm.location': None,
        'com.qualcomm.qti.qccauthmgr': None,
        'com.qualcomm.qti.qdma': None,
        'com.qualcomm.qti.qms.service.connectionsecurity': None,
        'com.qualcomm.qti.qms.service.trustzoneaccess': None,
        'com.qualcomm.qti.simcontacts': None
    }

    OVERWRITE_LIST = {
        'com.google.android.gms',  # empty list for not removing original data. PS: remove will cause serious issues
    }

    def __init__(self, context: Context):
        super().__init__(BackupAppDataAction.KEY, BackupAppDataAction.DISPLAY_NAME, context)

    def _on_backup(self, progress: Union[None, Callable[[int, int, str], None]]):
        device = self.wait_for_boot_up_ready()
        with device.shell() as shell:
            pkg_infos = Util.dump_pkg_infos(shell)
            user_list = pkg_infos.keys()

            # start user
            Action.start_users(user_list, shell=shell)
            Action.wait_for_filesystem_ready(user_list, shell=shell)

            # search for backup items
            backup_paths = shell.find('/storage/emulated/*/Android/*', 1, 1)\
                + shell.find('/data/user*/*/', 1, 1)\
                + shell.find('/data/misc/profiles/cur/*', 1, 1)
            backup_items = [BackupAppDataAction.BackupItem(path, self.data_dir(), pkg_infos) for path in backup_paths]
            backup_items = list(filter(lambda ii: ii.pkg_info is not None and ii.exclude is not None,
                                       backup_items))  # remove non pkg items & exclude item

            # group by pkg
            backup_dict: Dict[str, List[BackupAppDataAction.BackupItem]] = dict()
            for item in backup_items:
                item_list: List[BackupAppDataAction.BackupItem]
                if item.pkg in backup_dict:
                    item_list = backup_dict[item.pkg]
                else:
                    backup_dict[item.pkg] = item_list = []
                item_list.append(item)

            # do backup
            prog = ProgressHelper(progress, len(backup_dict))
            Log.d('stop android before backup app data')
            shell.android_stop()  # stop android first

            for pkg in backup_dict:
                prog.on_message(pkg).inc()  # progress

                item_list = backup_dict[pkg]
                self._on_backup_items(item_list, device, shell)
                Log.d('backup: %s at %s' % (pkg, [item.device_path for item in item_list]))

            Log.d('start android after backup app data')
            shell.android_start()  # start android

            prog.on_message('Done [%s]' % self.duration_time_str())
            Log.i('Backup success for %d app data' % len(backup_dict))

    def _on_restore(self, progress: Union[None, Callable[[int, int, str], None]]):
        device = self.wait_for_boot_up_ready()
        with device.shell() as shell:
            android_version = self.android_version()
            pkg_infos = Util.dump_pkg_infos(shell)
            user_list = pkg_infos.keys()

            # start user
            Action.start_users(user_list, shell=shell)
            Action.wait_for_filesystem_ready(user_list, shell=shell)

            # search for restore items
            restore_paths = [os.path.join(self.data_dir(), filename) for filename in os.listdir(self.data_dir())]
            restore_paths = filter(lambda p: not os.path.isdir(p), restore_paths)  # remove dirs
            restore_items = [BackupAppDataAction.RestoreItem(path, pkg_infos, android_version)
                             for path in restore_paths]
            restore_items = list(filter(lambda ii: ii.pkg_info is not None, restore_items))  # remove non pkg items

            # group by pkg
            restore_dict: Dict[str, List[BackupAppDataAction.RestoreItem]] = dict()
            for item in restore_items:
                item_list: List[BackupAppDataAction.RestoreItem]
                if item.pkg in restore_dict:
                    item_list = restore_dict[item.pkg]
                else:
                    restore_dict[item.pkg] = item_list = []
                item_list.append(item)

            # move com.android.launcher3 to last
            restore_dict_keys = list(restore_dict.keys())
            if 'com.android.launcher3' in restore_dict_keys:
                restore_dict_keys.remove('com.android.launcher3')
                restore_dict_keys.append('com.android.launcher3')

            # do restore
            prog = ProgressHelper(progress, len(restore_dict))
            Log.d('stop android before restoring app data')
            shell.android_stop()  # stop android first

            for pkg in restore_dict_keys:
                prog.on_message(pkg).inc()  # progress
                item_list = restore_dict[pkg]

                # restore files
                self._on_restore_items(item_list, device, shell)
                Log.d('restored: %s at %s' % (pkg, [item.device_path for item in item_list]))

            Log.d('start android after restoring app data')
            shell.android_start()  # start android

            prog.on_message('Done [%s]' % self.duration_time_str())

            Log.i('Restore success for %d app data' % len(restore_dict))

    def _on_backup_items(self, backup_items, device: AdbDevice, shell: AdbShell):
        if isinstance(backup_items, BackupAppDataAction.BackupItem):
            item = backup_items
            if item.pkg == 'com.android.launcher3':
                self._on_backup_launcher(item, device, shell)
            else:
                self._on_backup_generic(item, device, shell)
        elif isinstance(backup_items, Iterable):
            for item in backup_items:
                self._on_backup_items(item, device, shell)

    def _on_backup_generic(self, backup_item, device: AdbDevice, shell: AdbShell):
        item: BackupAppDataAction.BackupItem = backup_item
        self.backup_path(device_path=item.device_path,
                         host_dir=self.data_dir(),
                         device=device,
                         shell=shell,
                         compress=True,
                         compress_include=None,
                         compress_exclude=item.exclude)

    def _on_backup_launcher(self, backup_item, device: AdbDevice, shell: AdbShell):
        item: BackupAppDataAction.BackupItem = backup_item
        self._on_backup_generic(item, device, shell)

        if not item.device_path.startswith('/data/user/'):
            return  # only restore ext files when item under /data/user/

        host_dir = os.path.join(self.data_dir(), 'ext', item.pkg)
        if not os.path.exists(host_dir):
            os.makedirs(host_dir)

        device_path = '/data/system/users/%d/appwidgets.xml' % item.user_id
        if shell.path_exists(device_path):
            Log.d('backup launcher ext: %s' % device_path)
            self.backup_path(device_path=device_path,
                             host_dir=host_dir,
                             device=device,
                             shell=shell,
                             compress=False)

        device_path = '/data/system_ce/%d/shortcut_service' % item.user_id
        if shell.path_exists(device_path):
            Log.d('backup launcher ext: %s' % device_path)
            self.backup_path(device_path=device_path,
                             host_dir=host_dir,
                             device=device,
                             shell=shell,
                             compress=True)

    def _on_restore_items(self, restore_items, device: AdbDevice, shell: AdbShell):
        if isinstance(restore_items, BackupAppDataAction.RestoreItem):
            item = restore_items
            if item.pkg == 'com.android.launcher3':
                self._on_restore_launcher(item, device, shell)
            else:
                self._on_restore_generic(item, device, shell)
        elif isinstance(restore_items, Iterable):
            for item in restore_items:
                self._on_restore_items(item, device, shell)

    def _on_restore_generic(self, restore_item, device: AdbDevice, shell: AdbShell):
        item: BackupAppDataAction.RestoreItem = restore_item
        self.restore_path(host_path=item.host_path,
                          device=device,
                          shell=shell,
                          compress=True,
                          clear_copy=item.clear_copy)

        # chown
        shell.chown_on_exist(item.chown_u, item.chown_g, item.device_path)

        # chown for special folders
        if not item.uid.is_system_uid:
            owner = item.uid.name
            group = '%s_cache' % owner
            for cache_dir in {'cache', 'code_cache'}:
                shell.chown_on_exist(owner,
                                     group,
                                     '%s/%s' % (item.device_path, cache_dir))

    def _on_restore_launcher(self, restore_item, device: AdbDevice, shell: AdbShell):
        item: BackupAppDataAction.RestoreItem = restore_item
        self._on_restore_generic(item, device, shell)

        if not item.device_path.startswith('/data/user/'):
            return  # only restore ext files when item under /data/user/

        device_path = '/data/system/users/%d/appwidgets.xml' % item.user_id
        host_path = Action.device_path_to_host_path(device_path,
                                                    os.path.join(self.data_dir(), 'ext', item.pkg),
                                                    compress=False)
        if os.path.exists(host_path):
            Log.d('restore launcher ext: %s' % device_path)
            self.restore_path(host_path=host_path,
                              device=device,
                              shell=shell,
                              compress=False,
                              clear_copy=False)

            # chown
            shell.chown_on_exist('system', 'system', device_path)

            # chmod
            shell.chmod('600', device_path, recursive=False)

        device_path = '/data/system_ce/%d/shortcut_service' % item.user_id
        host_path = Action.device_path_to_host_path(device_path,
                                                    os.path.join(self.data_dir(), 'ext', item.pkg),
                                                    compress=True)
        if os.path.exists(host_path):
            Log.d('restore launcher ext: %s' % device_path)
            self.restore_path(host_path=host_path,
                              device=device,
                              shell=shell,
                              compress=True,
                              clear_copy=False)

    class BackupItem(object):
        PATTERNS = {
            re.compile(r'/storage/emulated/(\d+)/Android/(data|media)/(.+)'):
                lambda m: (int(m.group(1)), m.group(3)),
            re.compile(r'/data/(user|user_\w+)/(\d+)/(.+)'):
                lambda m: (int(m.group(2)), m.group(3)),
            re.compile(r'/data/misc/profiles/cur/(\d+)/(.+)'):
                lambda m: (int(m.group(1)), m.group(2)),
            # re.compile(r'/data/(misc_de|misc_ce)/(\d+)/apexdata/(.+)'):
            #     lambda m: (int(m.group(2)), m.group(3))
        }

        def __init__(self, device_path: str, host_dir: str,
                     pkg_infos: Dict[int, Dict[str, PmListPkgInfo]]):
            self.device_path = device_path
            self.host_dir = host_dir
            self.user_id = None
            self.pkg = None

            for pattern in BackupAppDataAction.BackupItem.PATTERNS:
                matched = pattern.fullmatch(self.device_path)
                if matched:
                    self.user_id, self.pkg = BackupAppDataAction.BackupItem.PATTERNS[pattern](matched)

            if self.user_id in pkg_infos:
                if self.pkg in pkg_infos[self.user_id]:
                    self.pkg_info = pkg_infos[self.user_id][self.pkg]
                else:
                    self.pkg_info = None
            else:
                self.pkg_info = None
            self.exclude = BackupAppDataAction.EXCLUDE_LIST[self.pkg]\
                if self.pkg in BackupAppDataAction.EXCLUDE_LIST\
                else ['cache', 'code_cache']

        def __str__(self):
            return self.device_path

    class RestoreItem(object):
        PATTERNS = {
            re.compile(r'-storage-emulated-(\d+)-Android-data/(.+).tar.gz'):
                lambda m, v: (int(m.group(1)), m.group(3), None, 'sdcard_rw'),
            re.compile(r'-storage-emulated-(\d+)-Android-media/(.+).tar.gz'):
                lambda m, v: (int(m.group(1)), m.group(3), None, 'everybody' if v >= 12 else 'sdcard_rw'),
            re.compile(r'-data-(user|user_\w+)-(\d+)-(.+).tar.gz'):
                lambda m, v: (int(m.group(2)), m.group(3), None, None),
            re.compile(r'-data-misc-profiles-cur-(\d+)-(.+).tar.gz'):
                lambda m, v: (int(m.group(1)), m.group(2), None, None),
            # re.compile(r'-data-(misc_de|misc_ce)-(\d+)-apexdata-(.+).tar.gz'):
            #     lambda m, v: (int(m.group(2)), m.group(3))
        }

        def __init__(self, host_path: str,
                     pkg_infos: Dict[int, Dict[str, PmListPkgInfo]],
                     android_version: int):
            self.host_path = host_path
            self.device_path = Action.host_path_to_device_path(host_path, compress=True)
            self.user_id = None
            self.pkg = None
            self.chown_u = None
            self.chown_g = None

            for pattern in BackupAppDataAction.RestoreItem.PATTERNS:
                matched = pattern.fullmatch(os.path.basename(self.host_path))
                if matched:
                    self.user_id, self.pkg, self.chown_u, self.chown_g\
                        = BackupAppDataAction.RestoreItem.PATTERNS[pattern](matched, android_version)

            if self.user_id in pkg_infos:
                if self.pkg in pkg_infos[self.user_id]:
                    self.pkg_info = pkg_infos[self.user_id][self.pkg]
                else:
                    self.pkg_info = None
            else:
                self.pkg_info = None
            self.clear_copy = self.pkg not in BackupAppDataAction.OVERWRITE_LIST
            self.uid = BackupAppDataAction.Uid(self.pkg_info.uid) if self.pkg_info is not None else None

            if self.uid:
                if self.chown_u is None:
                    self.chown_u = str(self.uid.uid)
                if self.chown_g is None:
                    self.chown_g = str(self.uid.uid)

        def __str__(self):
            return self.device_path

    class Uid(object):
        def __init__(self, uid: Union[int, str]):
            self.uid = uid if isinstance(uid, int) else int(uid)

            self.user_id = self.uid // 100000
            self.pkg_uid = self.uid % 100000

            self.is_system_uid = self.pkg_uid < 10000
            self.name = None if self.is_system_uid else 'u%d_a%d' % (self.user_id, self.pkg_uid - 10000)

        def __str__(self):
            return str(self.uid)
